package com.gitee.wsl.platform.thread.task.api

import com.gitee.wsl.platform.concurrent.api.Cancellable
import com.gitee.wsl.platform.concurrent.api.MutableCancellable
import com.gitee.wsl.platform.exception.logOrRethrow
import com.gitee.wsl.platform.exception.rethrowIfFatal
import com.gitee.wsl.platform.thread.task.TaskCancellationException
import com.gitee.wsl.platform.thread.task.TaskExecutor
import com.gitee.wsl.platform.thread.task.TaskLocalContext.signalTheStartOfBlockingCall
import com.gitee.wsl.platform.thread.task.bean.CancellableFuture
import com.gitee.wsl.platform.thread.task.bean.Outcome
import com.gitee.wsl.platform.thread.task.callback.AsyncContinuationCallback
import com.gitee.wsl.platform.thread.task.callback.AsyncContinuationFun
import com.gitee.wsl.platform.thread.task.callback.CancellableContinuation
import com.gitee.wsl.platform.thread.task.callback.CompletionCallback
import com.gitee.wsl.platform.thread.task.callback.Continuation
import com.gitee.wsl.platform.thread.task.collection.ImmutableQueue
import org.jetbrains.annotations.ApiStatus
import org.jetbrains.annotations.Blocking
import org.jetbrains.annotations.NonBlocking
import java.time.Duration
import java.util.*
import java.util.concurrent.*
import java.util.concurrent.atomic.AtomicReference
import java.util.concurrent.locks.AbstractQueuedSynchronizer
import java.util.function.Function


/**
 * A `Fiber` is a task that was started concurrently, and that can
 * be joined or cancelled.
 *
 *
 * Fibers are the equivalent of threads, except they are higher-level and
 * library-managed, instead of being intrinsic to the JVM.
 *
 * @param <T> is the result of the fiber, if successful.
</T> */
interface Fiber<T> : Cancellable {
    @Throws(
        ExecutionException::class,
        TaskCancellationException::class,
        NotCompletedException::class
    )
    @NonBlocking
    fun getResultOrThrow(): T

    /**
     * Waits until the fiber completes, and then runs the given callback to
     * signal its completion.
     *
     *
     * Completion includes cancellation. Triggering [.cancel] before
     * `joinAsync` will cause the fiber to get cancelled, and then the
     * "join" back-pressures on cancellation.
     *
     * @param onComplete is the callback to run when the fiber completes
     * (successfully, or with failure, or cancellation)
     *
     * @return a [Cancellable] that can be used to unregister the callback,
     * in case the caller is no longer interested in the result. Note this
     * does not cancel the fiber itself.
     */
    @NonBlocking
    fun joinAsync(onComplete: Runnable): Cancellable

    /**
     * Waits until the fiber completes, and then runs the given callback
     * to signal its completion.
     *
     *
     * This method can be executed as many times as necessary, with the
     * result of the `Fiber` being memoized. It can also be executed
     * after the fiber has completed, in which case the callback will be
     * executed immediately.
     *
     * @param callback will be called with the result when the fiber completes.
     *
     * @return a [Cancellable] that can be used to unregister the callback,
     * in case the caller is no longer interested in the result. Note this
     * does not cancel the fiber itself.
     */
    @NonBlocking
    fun awaitAsync(callback: CompletionCallback<in T>): Cancellable {
        return joinAsync(Runnable {
            try {
                val result = this.getResultOrThrow()
                callback.onSuccess(result)
            } catch (e: ExecutionException) {
                callback.onFailure(Objects.requireNonNullElse<Throwable>(e.cause, e))
            } catch (e: TaskCancellationException) {
                callback.onCancellation()
            } catch (e: Throwable) {
                e.rethrowIfFatal()
                callback.onFailure(e)
            }
        })
    }

    /**
     * Blocks the current thread until the fiber completes, or until
     * the timeout is reached.
     *
     *
     * This method does not return the outcome of the fiber. To check
     * the outcome, use [outcome].
     *
     * @throws InterruptedException if the current thread is interrupted, which
     * will just stop waiting for the fiber, but will not cancel the running
     * task.
     * @throws TimeoutException     if the timeout is reached before the fiber
     * completes.
     */
    @Blocking
    @Throws(InterruptedException::class, TimeoutException::class)
    fun joinBlockingTimed(timeout: Duration) {
        val latch = AwaitSignal()
        val token = joinAsync(Runnable { latch.signal() })
        try {
            latch.await(timeout.toMillis())
        } catch (e: InterruptedException) {
            token.cancel()
            throw e
        } catch (e: TimeoutException) {
            token.cancel()
            throw e
        }
    }

    /**
     * Blocks the current thread until the fiber completes, then returns
     * the result of the fiber.
     *
     * @param timeout is the maximum time to wait for the fiber to complete,
     * before throwing a [TimeoutException].
     * @return the result of the fiber, if successful.
     * @throws InterruptedException      if the current thread is interrupted, which
     * will just stop waiting for the fiber, but will not
     * cancel the running task.
     * @throws TimeoutException          if the timeout is reached before the fiber completes.
     * @throws TaskCancellationException if the fiber was cancelled concurrently.
     * @throws ExecutionException        if the task failed with an exception.
     */
    @Blocking
    @Throws(
        InterruptedException::class,
        TimeoutException::class,
        TaskCancellationException::class,
        ExecutionException::class
    )
    fun awaitBlockingTimed(timeout: Duration): T {
        joinBlockingTimed(timeout)
        try {
            return this.getResultOrThrow()
        } catch (e: NotCompletedException) {
            throw IllegalStateException(e)
        }
    }

    /**
     * Blocks the current thread until the fiber completes.
     *
     *
     * This method does not return the outcome of the fiber. To check
     * the outcome, use [outcome].
     *
     * @throws InterruptedException if the current thread is interrupted, which
     * will just stop waiting for the fiber, but will not cancel the running
     * task.
     */
    @Blocking
    @Throws(InterruptedException::class)
    fun joinBlocking() {
        val latch = AwaitSignal()
        val token = joinAsync(Runnable { latch.signal() })
        try {
            latch.await()
        } finally {
            token.cancel()
        }
    }

    /**
     * Blocks the current thread until the fiber completes.
     *
     *
     * Version of [.joinBlocking] that ignores thread interruptions.
     * This is most useful after cancelling a fiber, as it ensures that
     * processing will back-pressure on the fiber's completion.
     *
     *
     * **WARNING:** This method guarantees that upon its return
     * the fiber is completed, however, it still throws [InterruptedException]
     * because it can't swallow interruptions.
     *
     *
     * Sample:
     * <pre>`final var fiber = Task
     * .fromBlockingIO(() -> {
     * Thread.sleep(10000);
     * })
     * .runFiber();
     * // ...
     * fiber.cancel();
     * fiber.joinBlockingUninterruptible();
    `</pre> *
     */
    @Blocking
    @Throws(InterruptedException::class)
    fun joinBlockingUninterruptible() {
        var wasInterrupted = Thread.interrupted()
        while (true) {
            try {
                joinBlocking()
                break
            } catch (e: InterruptedException) {
                wasInterrupted = true
            }
        }
        if (wasInterrupted) {
            throw InterruptedException(
                "Thread was interrupted in #joinBlockingUninterruptible"
            )
        }
    }

    /**
     * Blocks the current thread until the fiber completes, then returns the
     * result of the fiber.
     *
     * @throws InterruptedException      if the current thread is interrupted, which
     * will just stop waiting for the fiber, but will not
     * cancel the running task.
     * @throws TaskCancellationException if the fiber was cancelled concurrently.
     * @throws ExecutionException        if the task failed with an exception.
     */
    @Blocking
    @Throws(InterruptedException::class, TaskCancellationException::class, ExecutionException::class)
    fun awaitBlocking(): T {
        joinBlocking()
        try {
            return this.getResultOrThrow()
        } catch (e: NotCompletedException) {
            throw IllegalStateException(e)
        }
    }

    /**
     * Cancels the fiber, which will eventually stop the running fiber (if
     * it's still running), completing it via "cancellation".
     *
     *
     * This manifests either in a [TaskCancellationException] being
     * thrown by [.getResultOrThrow], or in the
     * [CompletionCallback.onCancellation] callback being triggered.
     */
    @NonBlocking
    override fun cancel()

    /**
     * Returns a [CancellableFuture] that can be used to join the fiber
     * asynchronously, or to cancel it.
     *
     * @return a [CancellableFuture] that can be used to asynchronously
     * process its result, or to give up on listening for the result. Note that
     * cancelling the returned future does not cancel the fiber itself, for
     * that you need to call [Fiber.cancel].
     */
    @NonBlocking
    fun joinAsync(): CancellableFuture<Void> {
        val future = CompletableFuture<Void>()
        val token = joinAsync(Runnable { future.complete(null) })
        val cRef = Cancellable {
            try {
                token.cancel()
            } finally {
                future.cancel(false)
            }
        }
        return CancellableFuture(future, cRef)
    }

    /**
     * Overload of [.awaitAsync].
     *
     * @return a [CancellableFuture] that can be used to asynchronously
     * process its result or to give up on listening for the result. Note that
     * cancelling the returned future does not cancel the fiber itself, for
     * that you need to call [Fiber.cancel].
     */
    @NonBlocking
    fun awaitAsync(): CancellableFuture<T> {
        val f = joinAsync()
        return f.transform(Function { it: CompletableFuture<out Void> ->
            it.thenApply { v ->
                try {
                    return@thenApply this.getResultOrThrow()
                } catch (e: ExecutionException) {
                    throw CompletionException(e.cause)
                } catch (e: Throwable) {
                    e.rethrowIfFatal()
                    throw CompletionException(e)
                }
            }
        })
    }

    /**
     * Thrown in case [.getResultOrThrow] is called before the fiber
     * completes (i.e., before one of the "join" methods return).
     */
    class NotCompletedException : Exception("Fiber is not completed")
}

/**
 * INTERNAL API.
 *
 *
 * **INTERNAL API:** Internal apis are subject to change or removal
 * without any notice. When code depends on internal APIs, it is subject to
 * breakage between minor version updates.
 */
@ApiStatus.Internal
internal class ExecutedFiber<T> private constructor(
    private val executor: TaskExecutor
) : Fiber<T> {
    private val continuation: Continuation<T>
    private val cancellableRef: MutableCancellable
    private val stateRef: AtomicReference<State<T>>

    init {
        this.cancellableRef = MutableCancellable(Cancellable { this.fiberCancel() })
        this.stateRef = AtomicReference<State<T>>(State.start<T>())
        this.continuation = CancellableContinuation(
            executor,
            AsyncContinuationCallback<T>(
                FiberCallback<T>(executor, stateRef),
                executor
            ),
            cancellableRef
        )
    }

    private fun fiberCancel() {
        while (true) {
            val current = stateRef.get()
            if (current is State.Active<*>) {
                if (stateRef.compareAndSet(current, State.Cancelled(current.listeners))) {
                    return
                }
            } else {
                return
            }
        }
    }

    @Throws(ExecutionException::class, TaskCancellationException::class, Fiber.NotCompletedException::class)
    override fun getResultOrThrow(): T {
        val current = stateRef.get()
        if (current is State.Completed<*>) {
            return (current as State.Completed<T>).outcome.getOrThrow()
        } else {
            throw Fiber.NotCompletedException()
        }
    }

    override fun joinAsync(onComplete: Runnable): Cancellable {
        while (true) {
            val current = stateRef.get()
            if (current is State.Active<*> || current is State.Cancelled<*>) {
                val update = current.addListener(onComplete)
                if (stateRef.compareAndSet(current, update)) {
                    return removeListenerCancellable(onComplete)
                }
            } else {
                executor.resumeOnExecutor(onComplete)
                return Cancellable.EMPTY
            }
        }
    }

    override fun cancel() {
        this.cancellableRef.cancel()
    }

    private fun removeListenerCancellable(listener: Runnable): Cancellable {
        return Cancellable {
            while (true) {
                val current = stateRef.get()
                if (current is State.Active<*> || current is State.Cancelled<*>) {
                    val update = current.removeListener(listener)
                    if (stateRef.compareAndSet(current, update)) {
                        return@Cancellable
                    }
                } else {
                    return@Cancellable
                }
            }
        }
    }

    internal interface State<T> {
        class Active<T>(
            val listeners: ImmutableQueue<Runnable>,
        ) : State<T>

        class Cancelled<T>(
            val listeners: ImmutableQueue<Runnable>,
        ) : State<T>

        @JvmRecord
        data class Completed<T>(
            val outcome: Outcome<T>,
        ) : State<T>

        fun triggerListeners(executor: TaskExecutor) {
            if (this is Active<T>) {
                for (listener in listeners) {
                    executor.resumeOnExecutor(listener)
                }
            } else if (this is Cancelled<T>) {
                for (listener in listeners) {
                    executor.resumeOnExecutor(listener)
                }
            }
        }

        fun addListener(listener: Runnable): State<T> {
            if (this is Active<T>) {
                val newQueue = listeners.enqueue(listener)
                return Active(newQueue)
            } else if (this is Cancelled<T>) {
                val newQueue = listeners.enqueue(listener)
                return Cancelled<T>(newQueue)
            } else {
                return this
            }
        }

        fun removeListener(listener: Runnable): State<T> {
            if (this is Active<T>) {
                val newQueue = listeners.filter({ l -> l !== listener })
                return Active<T>(newQueue)
            } else if (this is Cancelled<T>) {
                val newQueue = listeners.filter({ l -> l !== listener })
                return Cancelled<T>(newQueue)
            } else {
                return this
            }
        }

        companion object {
            fun <T> start(): State<T> {
                return Active(ImmutableQueue.empty())
            }
        }
    }

    internal class FiberCallback<T>(
        private val executor: TaskExecutor,
        private val stateRef: AtomicReference<State<T>>,
    ) : CompletionCallback<T> {
        override fun onSuccess(value: T) {
            onOutcome(Outcome.success(value))
        }

        override fun onFailure(e: Throwable) {
            onOutcome(Outcome.failure(e))
        }

        override fun onCancellation() {
            onOutcome(Outcome.cancellation())
        }

        public override fun onOutcome(outcome: Outcome<T>) {
            while (true) {
                val current = stateRef.get()
                if (current is State.Active<*>) {
                    if (stateRef.compareAndSet(current, State.Completed(outcome))) {
                        current.triggerListeners(executor)
                        return
                    }
                } else if (current is State.Cancelled<*>) {
                    val update = State.Completed<T>(Outcome.cancellation())
                    if (stateRef.compareAndSet(current, update)) {
                        current.triggerListeners(executor)
                        return
                    }
                } else if (current is State.Completed<*>) {
                    if (outcome is Outcome.Failure<T>) {
                        outcome.exception.logOrRethrow()
                    }
                    return
                } else {
                    throw IllegalStateException("Invalid state: $current")
                }
            }
        }
    }

    companion object {
        fun <T> start(
            executor: Executor,
            createFun: AsyncContinuationFun<T>,
        ): Fiber<T> {
            val taskExecutor = TaskExecutor.from(executor)
            val fiber = ExecutedFiber<T>(taskExecutor)
            taskExecutor.execute({
                try {
                    createFun.invoke(fiber.continuation)
                } catch (e: Throwable) {
                    e.rethrowIfFatal()
                    fiber.continuation.onFailure(e)
                }
            })
            return fiber
        }
    }
}

/**
 * INTERNAL API.
 *
 *
 * **INTERNAL API:** Internal apis are subject to change or removal
 * without any notice. When code depends on internal APIs, it is subject to
 * breakage between minor version updates.
 */
@ApiStatus.Internal
internal class AwaitSignal : AbstractQueuedSynchronizer() {
    override fun tryAcquireShared(arg: Int): Int {
        return if (state != 0) 1 else -1
    }

    override fun tryReleaseShared(arg: Int): Boolean {
        state = 1
        return true
    }

    fun signal() {
        releaseShared(1)
    }

    @Throws(InterruptedException::class)
    fun await() {
        signalTheStartOfBlockingCall()
        acquireSharedInterruptibly(1)
    }

    @Throws(InterruptedException::class, TimeoutException::class)
    fun await(timeoutMillis: Long) {
        signalTheStartOfBlockingCall()
        if (!tryAcquireSharedNanos(1, TimeUnit.MILLISECONDS.toNanos(timeoutMillis))) {
            throw TimeoutException("Timed out after $timeoutMillis millis")
        }
    }
}
